from typing import Union, Dict, List, Optional, Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult, GenerationChunk, ChatGenerationChunk
from loguru import logger


class LLMCallbackHandler(BaseCallbackHandler):
    """Callback handler for streaming LLM responses."""

    def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], *, run_id: UUID,
                     parent_run_id: Optional[UUID] = None, tags: Optional[list[str]] = None,
                     metadata: Optional[dict[str, Any]] = None, **kwargs: Any) -> Any:
        logger.debug(f"on_llm_start:prompts={prompts}")

    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        # logger.debug(f"on_llm_new_token:chunk={chunk},token={token}")
        chunk = kwargs.get('chunk', None)
        reasoning_content = getattr(chunk.message, 'additional_kwargs', {}).get('reasoning_content'),
        logger.debug(f"on_llm_new_token:chunk={chunk},reasoning_content={reasoning_content}")

    def on_llm_end(self, response: LLMResult, *, run_id: UUID, parent_run_id: Optional[UUID] = None,
                   **kwargs: Any) -> Any:
        logger.debug(f"on_llm_end:response={response}")

    def on_llm_error(self, error: BaseException, *, run_id: UUID, parent_run_id: Optional[UUID] = None,
                     **kwargs: Any) -> Any:
        logger.debug(f"on_llm_error:error={error}")

    def on_tool_start(self, serialized: dict[str, Any], input_str: str, *, run_id: UUID,
                      parent_run_id: Optional[UUID] = None, tags: Optional[list[str]] = None,
                      metadata: Optional[dict[str, Any]] = None, inputs: Optional[dict[str, Any]] = None,
                      **kwargs: Any) -> Any:
        logger.debug(f"on_tool_start:input_str={input_str}")

    def on_tool_end(self, output: Any, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any) -> Any:
        logger.debug(f"on_tool_end:output={output}")

    def on_tool_error(self, error: BaseException, *, run_id: UUID, parent_run_id: Optional[UUID] = None,
                      **kwargs: Any) -> Any:
        logger.debug(f"on_tool_error:error={error}")
